#include "pendulum_demo.h"
#include <cmath>
#include <cstdlib>
#include <ctime>
#include <iostream>
#include <random>
#include <stdint.h> // 包含标准整数类型

Pendulum::Pendulum(mjModel *model, mjData *data) : m(model),
                                                   d(data),
                                                   noiseManager(sensor::NoiseManager::getInstance())
{
    initVisualization();
    initializeNoiseSettings();
}

Pendulum::~Pendulum()
{
    mjv_freeScene(&scn);
    mjr_freeContext(&con);
    if (window) {
        glfwDestroyWindow(window);
    }
    glfwTerminate();
}

void Pendulum::initializeNoiseSettings()
{
    // 初始化噪声管理器
    noiseManager.initialize(m);

    // 查找传感器ID
    int positionSensorId = mj_name2id(m, mjOBJ_SENSOR, "jointpos_sensor_1");
    int velocitySensorId = mj_name2id(m, mjOBJ_SENSOR, "jointvel_sensor_1");

    // 配置特定传感器的噪声
    if (positionSensorId != -1) {
        noiseManager.configureSensor(positionSensorId, sensor::NoiseType::GAUSSIAN, 0.2);
        LOG_INFO("Position sensor noise configured, ID: ", positionSensorId);
    } else {
        LOG_WARN("Position sensor not found");
    }

    if (velocitySensorId != -1) {
        noiseManager.configureSensor(velocitySensorId, sensor::NoiseType::GAUSSIAN, 0.5);
        LOG_INFO("Velocity sensor noise configured, ID: ", velocitySensorId);
    } else {
        LOG_WARN("Velocity sensor not found");
    }

    // 可以添加更多传感器配置...
}


void Pendulum::render()
{
    // 渲染数据准备
    mjrRect viewport = {0, 0, 0, 0};
    glfwGetFramebufferSize(window, &viewport.width, &viewport.height);

    // 更新和渲染场景
    mjv_updateScene(m, d, &opt, NULL, &cam, mjCAT_ALL, &scn);
    mjr_render(viewport, &scn, &con);
    glfwSwapBuffers(window);
    glfwPollEvents();
}

void Pendulum::updateCameraLookAt(double x)
{
    // 使用鼠标控制器更新相机位置
    mouseController.updateCameraLookAt(&cam, x);
}

/* ================================================================= */

void set_torque_control(const mjModel *m, mjData *d, int actuator_no, bool flag)
{
    if (flag)
        m->actuator_gainprm[10 * actuator_no + 0] = 1;
    else
        m->actuator_gainprm[10 * actuator_no + 0] = 0;
}


void set_position_servo(const mjModel *m, mjData *d, int actuator_no, double kp)
{
    // 设置指定执行器的增益参数为 kp
    m->actuator_gainprm[10 * actuator_no + 0] = kp;
    // 设置指定执行器的偏置参数为 -kp
    m->actuator_biasprm[10 * actuator_no + 1] = -kp;
}


void set_velocity_servo(const mjModel *m, mjData *d, int actuator_no, double kv)
{
    // 设置指定执行器的增益参数为 kv
    m->actuator_gainprm[10 * actuator_no + 0] = kv;
    // 设置指定执行器的偏置参数为 -kv
    m->actuator_biasprm[10 * actuator_no + 2] = -kv;
}


void applyPendulumControl(const mjModel *m, mjData *d)
{
    // 简单PD控制器
    // d->ctrl[0] = -10 * (d->sensordata[0] - 0) - 1 * d->sensordata[1];


    // 电机默认是开启的

    int actuator_no = 0;
    bool flag = true;
    set_torque_control(m, d, actuator_no, flag);
    d->ctrl[0] = -10 * (d->qpos[0] - 0) - 1 * d->qvel[0];


    actuator_no = 1;
    double kp = 0;

    set_position_servo(m, d, actuator_no, kp);
    // d->ctrl[1] = 0;
    // d->ctrl[1] = 1.57;
    // d->ctrl[1] = -1.57;


    actuator_no = 2;
    // double kv = 100;
    double kv = 0;
    set_velocity_servo(m, d, actuator_no, kv);
    // d->ctrl[2] = 0.2;

    LOG_INFO("Control: ", d->ctrl[0]);
    LOG_INFO("sensor_1: ", d->sensordata[0], " sensor_2: ", d->sensordata[1]);
}


/* ================================================================= */

// 修改控制回调函数部分
mjfGeneric Pendulum::myController()
{
    // 定义控制器回调函数
    static mjfGeneric callback = [](const mjModel *m, mjData *d) {
        // 获取噪声管理器单例
        sensor::NoiseManager &noiseManager = sensor::NoiseManager::getInstance();

        // 应用噪声到所有传感器数据
        // noiseManager.applyNoiseToAllSensors(m, d);

        // 可选：记录噪声应用情况（调试用）
        // noiseManager.logNoiseApplication(m, d);

        applyPendulumControl(m, d);
    };

    // 返回函数指针
    return callback;
}


void Pendulum::initVisualization()
{
    // 初始化GLFW和窗口
    mouseController.initializeGLFW(&window, "MuJoCo Inverted Pendulum", 1200, 900);

    // 初始化可视化组件
    mjv_defaultCamera(&cam);
    mjv_defaultOption(&opt);
    mjr_defaultContext(&con);

    // 检查输入有效性
    if (!m || !d) {
        std::cerr << "Invalid MuJoCo model or data pointers" << std::endl;
        throw std::invalid_argument("MuJoCo model or data is null");
    }

    // 初始化MuJoCo渲染器
    mjv_makeScene(m, &scn, 1000);
    mjr_makeContext(m, &con, mjFONTSCALE_150);

    // m->opt.gravity[2] = -1;
    // qpos is dim nqx1 = 7 x 1, 3 translations + 4 quaternions
    d->qpos[0] = 1.57;

    // 显示关节
    opt.flags[mjVIS_JOINT] = 1;

    // 设置相机参数
    mouseController.setupCamera(&cam, 90, -34, 5.6, 0.0, 0.0, 1.0);

    void *pointer = glfwGetWindowUserPointer(window);

    // 将相机对象指针关联到窗口
    glfwSetWindowUserPointer(window, pointer);
}
